package tuplestore

import (
	"errors"
	"sync"
)

// Original Index remains exactly the same
type Index struct {
	mu    sync.RWMutex
	items map[string]struct {
		segmentID string
		blockID   string
	}
}

func NewIndex() *Index {
	return &Index{
		items: make(map[string]struct {
			segmentID string
			blockID   string
		}),
	}
}

func (i *Index) LookupID(id string) (segmentID, blockID string, exists bool) {
	i.mu.RLock()
	defer i.mu.RUnlock()
	item, exists := i.items[id]
	return item.segmentID, item.blockID, exists
}

func (i *Index) Insert(id, segmentID, blockID string) {
	i.mu.Lock()
	defer i.mu.Unlock()
	i.items[id] = struct {
		segmentID string
		blockID   string
	}{segmentID, blockID}
}

// EnhancedIndex wraps the original Index and adds new functionality
type EnhancedIndex struct {
	*Index // Embed the original Index for backward compatibility

	// New functionality
	primaryIndex    *primaryKeyIndex
	decisionIndexes map[string]*decisionTreeIndex
	mu              sync.RWMutex
}

func NewEnhancedIndex() *EnhancedIndex {
	return &EnhancedIndex{
		Index:           NewIndex(), // Original index
		primaryIndex:    newPrimaryKeyIndex(),
		decisionIndexes: make(map[string]*decisionTreeIndex),
	}
}

// All original methods are automatically available via embedding
// We add new methods for enhanced functionality

// Primary key index implementation (unexported)
type primaryKeyIndex struct {
	mu         sync.RWMutex
	uniqueKeys map[string]bool
}

func newPrimaryKeyIndex() *primaryKeyIndex {
	return &primaryKeyIndex{
		uniqueKeys: make(map[string]bool),
	}
}

func (e *EnhancedIndex) EnablePrimaryKey() {
	e.primaryIndex = newPrimaryKeyIndex()
}

func (e *EnhancedIndex) InsertWithPrimaryKey(id, segmentID, blockID string) error {
	e.mu.Lock()
	defer e.mu.Unlock()

	if e.primaryIndex != nil {
		e.primaryIndex.mu.Lock()
		defer e.primaryIndex.mu.Unlock()
		if _, exists := e.primaryIndex.uniqueKeys[id]; exists {
			return errors.New("duplicate primary key")
		}
		e.primaryIndex.uniqueKeys[id] = true
	}

	// Still insert into the original index
	e.Index.Insert(id, segmentID, blockID)
	return nil
}

// Decision tree index implementation (unexported)
type decisionTreeIndex struct {
	root *decisionTreeNode
	mu   sync.RWMutex
}

type decisionTreeNode struct {
	key       string
	value     interface{}
	left      *decisionTreeNode
	right     *decisionTreeNode
	segmentID string
	blockID   string
}

func (e *EnhancedIndex) CreateDecisionIndex(fieldName string) {
	e.mu.Lock()
	defer e.mu.Unlock()

	if _, exists := e.decisionIndexes[fieldName]; !exists {
		e.decisionIndexes[fieldName] = &decisionTreeIndex{}
	}
}

func (e *EnhancedIndex) InsertWithIndex(id string, fields map[string]interface{}, segmentID, blockID string) error {
	// First insert into primary index if enabled
	if e.primaryIndex != nil {
		if err := e.InsertWithPrimaryKey(id, segmentID, blockID); err != nil {
			return err
		}
	} else {
		// Fall back to original insert
		e.Index.Insert(id, segmentID, blockID)
	}

	// Insert into decision tree indexes
	e.mu.Lock()
	defer e.mu.Unlock()

	for fieldName, value := range fields {
		if index, exists := e.decisionIndexes[fieldName]; exists {
			index.mu.Lock()
			// Simplified insertion for example
			// In real implementation, you'd properly balance the tree
			if index.root == nil {
				index.root = &decisionTreeNode{
					key:       id,
					value:     value,
					segmentID: segmentID,
					blockID:   blockID,
				}
			} else {
				// Basic BST insertion (would need balancing in production)
				current := index.root
				for {
					if id < current.key {
						if current.left == nil {
							current.left = &decisionTreeNode{
								key:       id,
								value:     value,
								segmentID: segmentID,
								blockID:   blockID,
							}
							break
						}
						current = current.left
					} else {
						if current.right == nil {
							current.right = &decisionTreeNode{
								key:       id,
								value:     value,
								segmentID: segmentID,
								blockID:   blockID,
							}
							break
						}
						current = current.right
					}
				}
			}
			index.mu.Unlock()
		}
	}

	return nil
}

func (e *EnhancedIndex) LookupByField(fieldName, key string) (segmentID, blockID string, exists bool) {
	e.mu.RLock()
	defer e.mu.RUnlock()

	if index, exists := e.decisionIndexes[fieldName]; exists {
		index.mu.RLock()
		defer index.mu.RUnlock()

		current := index.root
		for current != nil {
			if key == current.key {
				return current.segmentID, current.blockID, true
			} else if key < current.key {
				current = current.left
			} else {
				current = current.right
			}
		}
	}
	return "", "", false
}

func (e *EnhancedIndex) RangeQuery(fieldName, min, max string) []struct {
	Key       string
	Value     interface{}
	SegmentID string
	BlockID   string
} {
	e.mu.RLock()
	defer e.mu.RUnlock()

	var results []struct {
		Key       string
		Value     interface{}
		SegmentID string
		BlockID   string
	}

	if index, exists := e.decisionIndexes[fieldName]; exists {
		index.mu.RLock()
		defer index.mu.RUnlock()

		var traverse func(node *decisionTreeNode)
		traverse = func(node *decisionTreeNode) {
			if node == nil {
				return
			}

			if node.key >= min && node.key <= max {
				results = append(results, struct {
					Key       string
					Value     interface{}
					SegmentID string
					BlockID   string
				}{
					Key:       node.key,
					Value:     node.value,
					SegmentID: node.segmentID,
					BlockID:   node.blockID,
				})
			}

			if node.key > min {
				traverse(node.left)
			}
			if node.key < max {
				traverse(node.right)
			}
		}

		traverse(index.root)
	}

	return results
}
